import os
from PIL import Image
import torch
from torch.utils import data
import numpy as np
from torchvision import transforms as T
import torchvision
import cv2
import sys


class Dataset(data.Dataset):

    def __init__(self, root, data_list_file, phase='train', input_shape=(1, 128, 128)):
        self.phase = phase # 查看模式
        self.input_shape = input_shape

        with open(os.path.join(data_list_file), 'r') as fd: # 读取生成的txt文件，每一行为路径+类别标签
            imgs = fd.readlines()

        imgs = [os.path.join(root, img[:-1]) for img in imgs] # 组合每一张图片路径，img[:-1]为图片路径
        self.imgs = np.random.permutation(imgs) # 随机打乱imgs

        # normalize = T.Normalize(mean=[0.5, 0.5, 0.5],
        #                         std=[0.5, 0.5, 0.5])

        normalize = T.Normalize(mean=[0.5], std=[0.5]) # 归一化，

        if self.phase == 'train':
            self.transforms = T.Compose([ # 数据增强操作序列
                T.RandomCrop(self.input_shape[1:]),
                T.RandomHorizontalFlip(),
                T.ToTensor(),
                normalize
            ])
        else:
            self.transforms = T.Compose([
                T.CenterCrop(self.input_shape[1:]),
                T.ToTensor(),
                normalize
            ])

    def __getitem__(self, index):
        sample = self.imgs[index] # 索引获取该张图片的数据
        splits = sample.split() # 切分成列表，index0为路径，index1为标签
        img_path = splits[0]
        data = Image.open(img_path) # 打开图片获取像素值
        data = data.convert('L') # 转为灰度图像，每个像素用8个bit表示，0表示黑，255表示白，其他数字表示不同的灰度。
        data = self.transforms(data) # 图像增强
        label = np.int32(splits[1]) # 获取标签，并转化为int类型，方便后面进行损失计算以及acc
        return data.float(), label

    def __len__(self):
        return len(self.imgs)




if __name__ == '__main__': # 自行测试代码，train 不会走这里
    dataset = Dataset(root='/data/Datasets/fv/dataset_v1.1/dataset_mix_aligned_v1.1',
                      data_list_file='/data/Datasets/fv/dataset_v1.1/mix_20w.txt',
                      phase='test',
                      input_shape=(1, 128, 128))

    trainloader = data.DataLoader(dataset, batch_size=10)
    for i, (data, label) in enumerate(trainloader):
        # imgs, labels = data
        # print imgs.numpy().shape
        # print data.cpu().numpy()
        # if i == 0:
        img = torchvision.utils.make_grid(data).numpy()
        # print img.shape
        # print label.shape
        # chw -> hwc
        img = np.transpose(img, (1, 2, 0))
        # img *= np.array([0.229, 0.224, 0.225])
        # img += np.array([0.485, 0.456, 0.406])
        img += np.array([1, 1, 1])
        img *= 127.5
        img = img.astype(np.uint8)
        img = img[:, :, [2, 1, 0]]

        cv2.imshow('img', img)
        cv2.waitKey()
        # break
        # dst.decode_segmap(labels.numpy()[0], plot=True)
